Greptile Summary

This PR adds MuonOptimizer, a distributed Muon optimizer that applies SGD-momentum followed by distributed Newton-Schulz orthogonalization on tensor-parallel 2D parameter shards via an NCCL process group. The implementation is well-validated, mathematically correct relative to the full-matrix reference, and the closure/torch.enable_grad() interaction is handled properly.

• eps ≤ 0 is never validated: passing eps=0 causes silent division by zero inside _distributed_normalize_p2_ when a gradient norm is exactly zero (clamp_min(0) does not prevent sqrt(0) / 0).
• __del__ calling destroy() is a teardown ordering hazard: if dist.destroy_process_group() runs before the optimizer is garbage-collected, the finalizer will attempt to tear down a cuSolverMp context backed by a freed NCCL communicator.

Confidence Score: 5/5

Safe to merge; all findings are P2 suggestions with no blocking correctness or security issues.

All findings are P2 (style/hardening): missing eps validation, del teardown ordering, and import-time NUM_PROCS. The core optimizer math is correct, distributed normalization is equivalent to the full-matrix reference, and previously discussed issues (closure/enable_grad, global_shape scaling) are properly handled in this version.

transformer_engine/pytorch/optimizers/muon.py — eps validation gap and del teardown ordering.

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[step called] --> B{closure?}
    B -- yes --> C[enable_grad + call closure]
    B -- no --> D[iterate param groups]
    C --> D
    D --> E{p.grad is None?}
    E -- yes --> F[skip]
    E -- no --> G[resolve + validate partition_dim]
    G --> H{decoupled weight decay?}
    H -- yes --> I[p *= 1 - lr * wd]
    H -- no, wd != 0 --> J[grad += wd * p]
    I --> K[momentum_buffer.lerp_ grad]
    J --> K
    K --> L{nesterov?}
    L -- yes --> M[update = grad.lerp momentum_buffer momentum]
    L -- no --> N[update = momentum_buffer ref]
    M --> O[_orthogonalize update]
    N --> O
    O --> P[clone + maybe transpose]
    P --> Q[distributed_normalize_p2_ all_reduce norm]
    Q --> R[newton_schulz distributed]
    R --> S[maybe untranspose]
    S --> T[scale by get_muon_scale_factor * extra]
    T --> U[p += -lr * orth_update]
    U --> V[return loss]

_{Reviews (3): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

Having an optimizer with most code copied invites fragmentation.

The idea was to give something to users, who use TE, but not Megatron-LM. By fragmentation you mean that we want to encourage everyone to use Megatron-LM? Or that the optimizer being relatively thin thing on top of newton_schulz call, and the users should have no trouble creating it themselves?

I don't think we gain anything by putting it into tests, since we already have tests for newton_schulz call. So we need to decide whether we want this PR, or should abandon it altogether. @cyanguwa

Having an optimizer with most code copied invites fragmentation.

The idea was to give something to users, who use TE, but not Megatron-LM. By fragmentation you mean that we want to encourage everyone to use Megatron-LM? Or that the optimizer being relatively thin thing on top of newton_schulz call, and the users should have no trouble creating it themselves?

I don't think we gain anything by putting it into tests, since we already have tests for newton_schulz call. So we need to decide whether we want this PR, or should abandon it altogether. @cyanguwa

Fragmentation means there will be different flavor of muon in emerging optimizer and TE, also a lot of copied code. TE can have stalled feature when emerging optimizer updates. Megatron-LM will always have its own version because there are implementation specific things need to be hooked together. For example, how QKV is implemetned, or fused swighlu.
For TE, I think an example of how to build a version of emerging optimizer use TE NS backend would be good to have. But providing optimizer (not fusion optimized version) confuses customers.
Having said that, I would love for TE to have a more optimized version. similar idea as fusedAdam, etc.